from elasticsearch7 import Elasticsearch, helpers
from nltk.stem import PorterStemmer
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import nltk
import re

import warnings
warnings.simplefilter("ignore")  # 屏蔽ES的一些warnings

nltk.download('punkt')  # 英文切词，词根，切句等方法
nltk.download('stopwords')  # 英文停用词库

def to_keywords(input_string):
    '''(英文)文本只保留关键字'''
    # 使用正则表达式替换所有非字母数字的字符为空格
    no_symbols = re.sub(r'[^a-zA-Z0-9\s]', ' ', input_string)
    word_tokens = word_tokenize(no_symbols)
    # 加载停用词表
    stop_words = set(stopwords.words('english'))
    ps = PorterStemmer()
    # 去停用词，取词根
    filtered_sentence = [ps.stem(w) for w in word_tokens if not w.lower() in stop_words]
    return ' '.join(filtered_sentence)

import os, time

# 以下代码需要 elasticsearch 数据库支持
# 引入配置文件
ELASTICSEARCH_BASE_URL = os.getenv('ELASTICSEARCH_BASE_URL')
ELASTICSEARCH_PASSWORD = os.getenv('ELASTICSEARCH_PASSWORD')
ELASTICSEARCH_USERNAME = os.getenv('ELASTICSEARCH_USERNAME')

# 1. 创建ES客户端
es = Elasticsearch(
    ELASTICSEARCH_BASE_URL,  # ES地址
    http_auth=(ELASTICSEARCH_USERNAME, ELASTICSEARCH_PASSWORD), # 用户名，密码
)

# 2. 定义索引名称
index_name = 'teacher_demo_index0'

# 3. 如果索引已存在，删除它（仅供演示，实际应用时不需要这步）
if es.indices.exists(index=index_name):  # 判断索引是否存在
    es.indices.delete(index=index_name)

# 4. 创建索引
es.indices.create(index=index_name)

from rag01 import extract_text_from_pdf
paragraphs = extract_text_from_pdf('llama2.pdf', min_line_length=10)

# 5. 灌库指令
actions = [
    {
        '_index': index_name,
        '_id': 1,
        '_source': {
            "keywords": to_keywords(para),
            "text": para
       }
    } for para in paragraphs
]

# 6. 文本灌库
helpers.bulk(es, actions)

# 灌库师异步的
time.sleep(2)

def search(query_string, top_n=3):
    # ES 的查询语句
    search_query = {
        "match": {
            "keywords": to_keywords(query_string)
        }
    }

    res = es.search(index=index_name, query=search_query, size=top_n)
    return [hit["_source"]["text"] for hit in res["hits"]["hits"]]

if __name__ == "__main__":
    # pass
    results = search("how many params does llama2 have?", 2)
    for r in results:
        print(r + "\n")